package com.example.airole.llm.service.impl;

import com.example.airole.llm.entity.ChatMessage;
import com.example.airole.llm.entity.ChatRequest;
import com.example.airole.llm.entity.ChatResponse;
import com.example.airole.common.entity.Role;
import com.example.airole.llm.service.LLMService;
import com.example.airole.common.mapper.RoleMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.*;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestTemplate;

import java.util.*;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Ollama LLM服务实现
 * 使用本地部署的Ollama服务
 */
@Slf4j
@Service("ollamaLLMService")
public class OllamaLLMService implements LLMService {
    
    @Value("${airole.llm.ollama.base-url:http://localhost:11434}")
    private String baseUrl;
    
    @Value("${airole.llm.ollama.model:qwen3:4b}")
    private String defaultModel;
    
    @Value("${airole.llm.ollama.enabled:true}")
    private boolean enabled;
    
    @Autowired
    private RoleMapper roleMapper;
    
    private final RestTemplate restTemplate = new RestTemplate();
    private final ExecutorService executorService = Executors.newFixedThreadPool(5);
    private final Map<String, ChatResponse> asyncResults = new ConcurrentHashMap<>();
    
    private static final String CHAT_ENDPOINT = "/api/chat";
    private static final String GENERATE_ENDPOINT = "/api/generate";
    private static final String MODELS_ENDPOINT = "/api/tags";
    
    @Override
    public ChatResponse chat(ChatRequest request) {
        if (!enabled) {
            return ChatResponse.builder()
                    .success(false)
                    .errorMessage("Ollama LLM服务未启用")
                    .build();
        }
        
        try {
            long startTime = System.currentTimeMillis();
            
            String model = request.getModel() != null ? request.getModel() : defaultModel;
            String roleId = request.getRoleId();
            
            // 构建聊天消息
            List<Map<String, Object>> messages = buildChatMessages(request, roleId);
            
            // 构建请求体
            Map<String, Object> requestBody = new HashMap<>();
            requestBody.put("model", model);
            requestBody.put("messages", messages);
            requestBody.put("stream", false);
            requestBody.put("options", buildChatOptions(request));
            
            // 发送请求
            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.APPLICATION_JSON);
            HttpEntity<Map<String, Object>> entity = new HttpEntity<>(requestBody, headers);
            
            String url = baseUrl + CHAT_ENDPOINT;
            ResponseEntity<Map> response = restTemplate.postForEntity(url, entity, Map.class);
            
            if (response.getStatusCode() == HttpStatus.OK && response.getBody() != null) {
                Map<String, Object> responseBody = response.getBody();
                
                long processingTime = System.currentTimeMillis() - startTime;
                
                String content = extractContentFromResponse(responseBody);
                Integer tokenUsage = extractTokenUsage(responseBody);
                
                return ChatResponse.builder()
                        .success(true)
                        .content(content)
                        .model(model)
                        .usageTokens(tokenUsage)
                        .processingTime(processingTime)
                        .build();
            } else {
                return ChatResponse.builder()
                        .success(false)
                        .errorMessage("LLM聊天失败，HTTP状态码: " + response.getStatusCode())
                        .build();
            }
            
        } catch (Exception e) {
            log.error("Ollama LLM聊天异常", e);
            return ChatResponse.builder()
                    .success(false)
                    .errorMessage("LLM聊天失败: " + e.getMessage())
                    .build();
        }
    }
    
    @Override
    public List<ChatResponse> chatStream(ChatRequest request) {
        if (!enabled) {
            return List.of(ChatResponse.builder()
                    .success(false)
                    .errorMessage("Ollama LLM服务未启用")
                    .build());
        }
        
        // 简化实现：流式聊天暂时返回单次响应
        // 实际项目中可以使用WebFlux的WebClient来实现真正的流式响应
        return List.of(chat(request));
    }
    
    @Override
    public String chatAsync(ChatRequest request) {
        String taskId = "llm_task_" + System.currentTimeMillis() + "_" + Thread.currentThread().getId();
        
        executorService.submit(() -> {
            try {
                ChatResponse response = chat(request);
                asyncResults.put(taskId, response);
            } catch (Exception e) {
                log.error("异步LLM聊天异常", e);
                asyncResults.put(taskId, ChatResponse.builder()
                        .success(false)
                        .errorMessage("异步LLM聊天失败: " + e.getMessage())
                        .build());
            }
        });
        
        return taskId;
    }
    
    @Override
    public ChatResponse getAsyncResult(String taskId) {
        return asyncResults.get(taskId);
    }
    
    @Override
    public boolean isServiceAvailable() {
        if (!enabled) {
            return false;
        }
        
        try {
            String url = baseUrl + MODELS_ENDPOINT;
            ResponseEntity<Map> response = restTemplate.getForEntity(url, Map.class);
            return response.getStatusCode() == HttpStatus.OK;
        } catch (Exception e) {
            log.warn("检查Ollama服务状态失败", e);
            return false;
        }
    }
    
    @Override
    public String[] getSupportedModels() {
        try {
            String url = baseUrl + MODELS_ENDPOINT;
            ResponseEntity<Map> response = restTemplate.getForEntity(url, Map.class);
            
            if (response.getStatusCode() == HttpStatus.OK && response.getBody() != null) {
                Map<String, Object> body = response.getBody();
                List<Map<String, Object>> models = (List<Map<String, Object>>) body.get("models");
                
                return models.stream()
                        .map(model -> (String) model.get("name"))
                        .toArray(String[]::new);
            }
        } catch (Exception e) {
            log.error("获取支持的模型列表失败", e);
        }
        
        return new String[]{defaultModel};
    }
    
    @Override
    public Role createRole(Role role) {
        if (role.getRoleId() == null || role.getRoleId().isEmpty()) {
            role.setRoleId("role_" + System.currentTimeMillis());
        }
        
        // 检查角色ID是否已存在
        if (roleMapper.countByRoleId(role.getRoleId()) > 0) {
            throw new RuntimeException("角色ID已存在: " + role.getRoleId());
        }
        
        // 设置默认值
        if (role.getCreateTime() == null) {
            role.setCreateTime(java.time.LocalDateTime.now());
        }
        if (role.getUpdateTime() == null) {
            role.setUpdateTime(java.time.LocalDateTime.now());
        }
        if (role.getUsageCount() == null) {
            role.setUsageCount(0L);
        }
        if (role.getIsPublic() == null) {
            role.setIsPublic(true);
        }
        
        int result = roleMapper.insert(role);
        if (result > 0) {
            return roleMapper.selectByRoleId(role.getRoleId());
        } else {
            throw new RuntimeException("创建角色失败");
        }
    }
    
    @Override
    public Role getRole(String roleId) {
        return roleMapper.selectByRoleId(roleId);
    }
    
    @Override
    public Role updateRole(Role role) {
        if (role.getRoleId() == null) {
            throw new RuntimeException("角色ID不能为空");
        }
        
        // 检查角色是否存在
        if (roleMapper.countByRoleId(role.getRoleId()) == 0) {
            throw new RuntimeException("角色不存在: " + role.getRoleId());
        }
        
        // 更新修改时间
        role.setUpdateTime(java.time.LocalDateTime.now());
        
        int result = roleMapper.updateByRoleId(role);
        if (result > 0) {
            return roleMapper.selectByRoleId(role.getRoleId());
        } else {
            throw new RuntimeException("更新角色失败");
        }
    }
    
    @Override
    public boolean deleteRole(String roleId) {
        if (roleId == null || roleId.isEmpty()) {
            return false;
        }
        
        int result = roleMapper.deleteByRoleId(roleId);
        return result > 0;
    }
    
    @Override
    public Role[] getRoles(String category, Boolean isPublic) {
        List<Role> roles = roleMapper.selectByCondition(category, isPublic);
        return roles.toArray(new Role[0]);
    }
    
    /**
     * 构建聊天消息
     */
    private List<Map<String, Object>> buildChatMessages(ChatRequest request, String roleId) {
        List<Map<String, Object>> messages = new ArrayList<>();
        
        // 添加系统消息（如果有角色设定）
        if (roleId != null) {
            Role role = getRole(roleId);
            if (role != null && role.getSystemPrompt() != null && !role.getSystemPrompt().isEmpty()) {
                Map<String, Object> systemMessage = new HashMap<>();
                systemMessage.put("role", "system");
                systemMessage.put("content", role.getSystemPrompt());
                messages.add(systemMessage);
            }
        }
        
        // 添加用户消息
        for (ChatMessage message : request.getMessages()) {
            Map<String, Object> msg = new HashMap<>();
            msg.put("role", message.getRole());
            msg.put("content", message.getContent());
            messages.add(msg);
        }
        
        return messages;
    }
    
    /**
     * 构建聊天选项
     */
    private Map<String, Object> buildChatOptions(ChatRequest request) {
        Map<String, Object> options = new HashMap<>();
        
        if (request.getTemperature() != null) {
            options.put("temperature", request.getTemperature());
        }
        
        if (request.getMaxTokens() != null) {
            options.put("num_predict", request.getMaxTokens());
        }
        
        if (request.getTopP() != null) {
            options.put("top_p", request.getTopP());
        }
        
        return options;
    }
    
    /**
     * 从响应中提取内容
     */
    private String extractContentFromResponse(Map<String, Object> responseBody) {
        try {
            Map<String, Object> message = (Map<String, Object>) responseBody.get("message");
            if (message != null) {
                return (String) message.get("content");
            }
        } catch (Exception e) {
            log.error("提取响应内容失败", e);
        }
        return "";
    }
    
    /**
     * 从响应中提取令牌使用量
     */
    private Integer extractTokenUsage(Map<String, Object> responseBody) {
        try {
            // 尝试直接获取 eval_count 作为 Integer
            Object evalCountObj = responseBody.get("eval_count");
            if (evalCountObj instanceof Integer) {
                return (Integer) evalCountObj;
            }
            
            // 如果 eval_count 是一个 Map，则从中获取 eval_count 字段
            if (evalCountObj instanceof Map) {
                Map<String, Object> evalCount = (Map<String, Object>) evalCountObj;
                Object countObj = evalCount.get("eval_count");
                if (countObj instanceof Integer) {
                    return (Integer) countObj;
                }
            }
            
            // 尝试其他可能的字段名
            Object promptEvalCountObj = responseBody.get("prompt_eval_count");
            if (promptEvalCountObj instanceof Integer) {
                return (Integer) promptEvalCountObj;
            }
            
        } catch (Exception e) {
            log.error("提取令牌使用量失败", e);
        }
        return 0;
    }
}